Skip to content

Add customization of redistribution build template #163

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
ybaturina merged 13 commits into from
Jan 9, 2026
Merged

Add customization of redistribution build template #163

ybaturina merged 13 commits into from
Jan 9, 2026

Conversation

@ybaturina
Copy link
Collaborator

@ybaturina ybaturina commented Jan 6, 2026 •
edited
Loading

When a custom source of redistribution is provided, we might need to have a custom build template too.

Example of customization:

load(
    "@cuda_redist_json//:distributions.bzl",
    "CUDA_REDISTRIBUTIONS",
)
load(
    "@rules_ml_toolchain//third_party/gpus/cuda/hermetic:cuda_redist_versions.bzl",
    "REDIST_VERSIONS_TO_BUILD_TEMPLATES",
)

CCCL_DIST_DICT = {
    "cuda_cccl": {
        "linux-x86_64": {
            "full_path": "https://github.com/NVIDIA/cccl/archive/0d328e06c9fc78a216ec70df4917f7230a9c77e3.tar.gz",
            "sha256": "c45dddfcebfc2d719e0c4cc6a874a4b50a751b90daba139699d3fc11708cf0ef",
            "strip_prefix": "cccl-0d328e06c9fc78a216ec70df4917f7230a9c77e3",
        },
        "linux-sbsa": {
            "full_path": "https://github.com/NVIDIA/cccl/archive/0d328e06c9fc78a216ec70df4917f7230a9c77e3.tar.gz",
            "sha256": "c45dddfcebfc2d719e0c4cc6a874a4b50a751b90daba139699d3fc11708cf0ef",
            "strip_prefix": "cccl-0d328e06c9fc78a216ec70df4917f7230a9c77e3",
        },
    },
}

CCCL_BUILD_TEMPLATES = {
    "cuda_cccl": {
        "repo_name": "cuda_cccl",
        "version_to_template": {
            "13": "@rules_ml_toolchain//third_party/gpus/cuda/hermetic:cuda_cccl_github.BUILD.tpl",
            "12": "@rules_ml_toolchain//third_party/gpus/cuda/hermetic:cuda_cccl_github.BUILD.tpl",
            "11": "@rules_ml_toolchain//third_party/gpus/cuda/hermetic:cuda_cccl_github.BUILD.tpl",
        },
        "local": {
            "source_dirs": ["include", "lib"],
            "version_to_template": {
                "13": "@rules_ml_toolchain//third_party/gpus/cuda/hermetic:cuda_cccl.BUILD.tpl",
                "12": "@rules_ml_toolchain//third_party/gpus/cuda/hermetic:cuda_cccl.BUILD.tpl",
                "11": "@rules_ml_toolchain//third_party/gpus/cuda/hermetic:cuda_cccl.BUILD.tpl",
            },
        },
    },
}

cuda_redist_init_repositories(
    cuda_redistributions = CUDA_REDISTRIBUTIONS | CCCL_DIST_DICT,
    redist_versions_to_build_templates = REDIST_VERSIONS_TO_BUILD_TEMPLATES | CCCL_BUILD_TEMPLATES,
)

In the scenario above cuda_cccl redistribution will be downloaded from Github, and the BUILD template will be used as in @rules_ml_toolchain//third_party/gpus/cuda/hermetic:cuda_cccl_github.BUILD.tpl.

For LOCAL_CUDA_PATH scenario, a template @rules_ml_toolchain//third_party/gpus/cuda/hermetic:cuda_cccl.BUILD.tpl will be used.

Another improvement is introducing strip_prefix on the level of individual redistribution URL, which fixes the issue when the root folder name in the archive is different from the file name.

@ybaturina ybaturina changed the title Add feature to customize build templates Add customization of redistribution build template Jan 6, 2026
@ybaturina ybaturina requested a review from vam-google January 6, 2026 23:18
@ybaturina ybaturina mentioned this pull request Jan 7, 2026
Closed
vam-google
vam-google approved these changes Jan 8, 2026
View reviewed changes
Copy link
Collaborator

@vam-google vam-google left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@ybaturina ybaturina merged commit 7e588e5 into google-ml-infra:main Jan 9, 2026
2 checks passed
@ybaturina ybaturina deleted the yb-custom-build branch January 9, 2026 18:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@vam-google vam-google vam-google approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@ybaturina @vam-google